from __future__ import with_statement
from khronos import *
from khronos.extras.components import *
from random import Random
from collections import deque
from contextlib import contextmanager
import re

class Input(object):
    def __init__(self, product, qty=1, pattern=False):
        self.product = product
        self.qty = qty
        self.pattern = re.compile(product) if pattern else None
        
    def __repr__(self):
        return "%d '%s'" % (self.qty, self.product)
        
    def __len__(self):
        return self.qty
        
    def match(self, product):
        if self.pattern is None:
            if self.product == product.name:
                return {}
            else:
                return None
        else:
            match = self.pattern.search(product.name)
            if match:
                return match.groupdict()
            else:
                return None
                
    def check_availability(self, available):
        if self.pattern is None:
            if self.product in available:
                return self.product
        else:
            for pname in available:
                if self.pattern.search(pname) is not None:
                    return pname
        return None
        
class Output(object):
    def __init__(self, bins, qty=1, mother=None):
        self.bins = bins
        self.qty = qty
        self.mother = mother
        
    def __repr__(self):
        return "%d [%s]%s" % (self.qty, ", ".join(str(bin) for bin in self.bins), 
                              "<%s>" % (self.mother.product) if self.mother else "")
        
    def link(self, distrs):
        for bin in self.bins:
            bin.distr = distrs[bin.distr]
            
    def generate(self, times):
        fractions = []
        fractions_total = 0.0
        for bin in self.bins:
            f = max(0, bin.distr())
            fractions.append((bin, f))
            fractions_total += f
            
        diffs = []
        total_qty = 0
        for bin, fraction in fractions:
            qty = fraction / fractions_total * times * self.qty
            int_qty = int(qty)
            diffs.append((qty - int_qty, int_qty, bin))
            total_qty += int_qty
            
        quantities = None
        remaining = self.qty * times - total_qty
        if remaining:
            quantities = []
            diffs.sort()
            while remaining:
                _, size, bin = diffs.pop()
                quantities.append((bin, size + 1))
                remaining -= 1
            quantities.extend([(bin, qty) for (_, qty, bin) in diffs])
        else:
            quantities = [(bin, qty) for (_, qty, bin) in diffs]
        return quantities
        
class Bin(object):
    def __init__(self, product, distr, inherits=False):
        self.product = product
        self.distr = distr
        self.inherits = inherits
        
    def __repr__(self):
        try:
            distr_name = self.distr.__name__
        except AttributeError:
            distr_name = self.distr
        return "%s: '%s'%s" % (distr_name, self.product, "*" if self.inherits else "")
        
class Process(object):
    def __init__(self, input, output, tt=None, tbf=None, tol=None):
        self.input = input
        self.output = output
        self.tt = tt
        self.tbf = tbf
        self.tol = tol
        
    def __repr__(self):
        return "%s >> %s" % (" + ".join(str(i) for i in self.input), 
                             " + ".join(str(o) for o in self.output))
                             
    def link(self, distrs):
        self.tt = distrs["tt"]
        self.tbf = distrs["tbf"]
        self.tol = distrs["tol"]
        for o in self.output:
            o.link(distrs)
            
    def check_availability(self, products):
        available = set(products)
        matched = []
        for item in self.input:
            match = item.check_availability(available)
            if match is None:
                return None
            matched.append(match)
            available.remove(match)
        return matched
        
    def match(self, input):
        """Matches a set of lots to the process' input definition. During the matching process, a 
        namespace is created for use in output regexp expansion, since output items can use 
        information from the input regexes. The couplings dict contains the bindings between the 
        items in the process' input definition and the provided lots. Additionally, a calculation 
        of the number of times the process may be applied to the input set is done."""
        provided_input = set(input)
        namespace = {}
        couplings = {}
        times = INF
        for item in self.input:
            matched = None
            for provided_item in provided_input:
                matched = item.match(provided_item.product)
                if matched is not None:
                    namespace.update(matched)
                    couplings[item] = provided_item
                    provided_input.remove(provided_item)
                    times = min(times, len(provided_item) // len(item))
                    break
            if matched is None:
                return None, None, None
        if times <= 0:
            raise Value_error("not enough qty to apply process to %s" % (input,))
        return couplings, namespace, times
        
    def apply(self, input):
        couplings, namespace, times = self.match(input)
        if times is None:
            return None
        else:
            return Applied_process(self, couplings, namespace, times)
            
class Applied_process(object):
    def __init__(self, process, couplings, namespace, times):
        self.process = process
        self.couplings = couplings
        self.namespace = namespace
        self.times = times
        
    def tool_time(self):
        return self.process.tt() * self.times
        
    def generate_fail(self):
        return self.process.tbf(), self.process.tol()
        
    def __call__(self, product_table):
        # decrease input lot size
        for input, lot in self.couplings.iteritems():
            lot.size -= self.times * len(input)
            
        # create output lots
        output_lots = []
        for output in self.process.output:
            mother_lot = None
            if output.mother is not None:
                mother_lot = self.couplings[output.mother]
                
            for bin, size in output.generate(self.times):
                product = product_table[bin.product % self.namespace]
                new_lot = Lot(product, size, mother=mother_lot)
                if bin.inherits:
                    new_lot.number, mother_lot.number = mother_lot.number, new_lot.number
                output_lots.append(new_lot)
                
        # filter out empty input lots
        remaining_input = [lot for lot in self.couplings.itervalues() if len(lot) > 0]
        return remaining_input, output_lots
        
class Lot(object):
    __count = 0
    def __init__(self, product, size, mother=None):
        self.product = product
        self.size = size
        self.mother = mother
        self.number = Lot.__count
        Lot.__count += 1
        
    def __repr__(self):
        return "%s(%s:%d, %d)" % (self.__class__.__name__, self.product, self.number, self.size)
        
    def __len__(self):
        return self.size
        
class Product(object):
    def __init__(self, name, route=None):
        self.name = name
        self.route = route
        
    def __repr__(self):
        return self.name
        
class Route(object):
    def __init__(self, name, sequence=()):
        self.name = name
        self.sequence = list(sequence)
        
    def __repr__(self):
        return "%s(%s, %s)" % (self.__class__.__name__, self.name, self.sequence)
        
    def __getitem__(self, idx):
        return self.sequence[idx]
        
    def __iter__(self):
        return iter(self.sequence)
        
    def after(self, operation):
        pos = self.sequence.index(operation)
        return self.sequence[pos + 1]
        
class Release(object):
    def __init__(self, product, lot_size, rate):
        self.product = product
        self.lot_size = lot_size
        self.rate = rate
        
    def mk_lot(self):
        return Lot(product=self.product, size=self.lot_size())
        
class Eq_setup(object):
    def __init__(self, name, conversions=None, processes=()):
        self.name = name
        self.conversions = conversions or {}
        self.processes = set(processes)
        self.local_slots = set([(self, p) for p in processes])
        
    def __str__(self):
        return self.name
        
    def add_process(self, proc):
        self.processes.add(proc)
        self.local_slots.add((self, proc))
        
    def slots(self):
        for listener in self.local_slots:
            yield listener
        for neighbor in self.conversions.iterkeys():
            for listener in neighbor.local_slots:
                yield listener
                
    def conversion_time(self, tgt_setup):
        return self.conversions[tgt_setup]()
        
class Eq_type(dict):
    def __init__(self, name):
        self.name = name
        
    def __str__(self):
        return self.name
        
    def add_setup(self, setup, *args, **kwargs):
        new_setup = None
        if isinstance(setup, EqSetup):
            new_setup = setup
            setup = setup.name
        else:
            if setup in self:
                raise Key_error("duplicate setup '%s' in equipment type '%s'" % (setup, self.name))
            new_setup = Eq_setup(setup, *args, **kwargs)
        self[setup] = new_setup
        return new_setup
        
    def set_conversions(self, **mapping):
        for src_setup, conversions in mapping.iteritems():
            src = self[src_setup]
            for dst_setup, time, in conversions.iteritems():
                dst = self[dst_setup]
                src.conversions[dst] = time
                
class Store(object):
    def __init__(self):
        self.content = {}
        self.size = 0
        
    def __str__(self):
        lines = []
        for pname, lots in sorted(self.content.iteritems()):
            lines.append("%s: %s" % (pname, lots))
        return "\n".join(lines)
        
    def __len__(self):
        return self.size
        
    def clear(self):
        self.content = {}
        self.size = 0
        
    def products(self):
        return self.content.keys()
        
    def add(self, lot):
        try:
            self.content[lot.product.name].append(lot)
        except KeyError:
            self.content[lot.product.name] = deque([lot])
        self.size += 1
        
    def get(self, products, verify=True):
        if verify and not self.has(products):
            raise Value_error("not all products are available in store")
        lots = []
        for product in products:
            l = self.content[product]
            lots.append(l.popleft())
            if not l:
                del self.content[product]
        self.size -= len(lots)
        return lots
        
    def has(self, products):
        for product in products:
            if product not in self.content:
                return False
        return True
        
class Skill(object):
    def __init__(self, func, time):
        self.func = func
        self.time = time
        self.name = func.__name__
        
    def __str__(self):
        return self.name
        
    def __call__(self, worker, *args, **kwargs):
        return self.func(worker, self.time(), *args, **kwargs)
        
class Skillset(dict):
    def __init__(self, name, *skills):
        self.name = name
        self.add(*skills)
        
    def __str__(self):
        return self.name
        
    def add(self, *skills):
        for skill in skills:
            self[skill.name] = skill
            
# ++++++++++++++++++++++++++++++++++++++++++++++++
class LotRelease(Composite):
    def initializer(self):
        for rdata in self.get_config():
            self.add(Builder(Call(self.release, rdata), 
                             period=rdata.rate, 
                             name=str(rdata.product)))
                             
    def release(self, rdata):
        self.send(rdata.mk_lot(), rdata.product.route[0])
        
    def finalizer(self):
        self.clear()
        
class LotDispatcher(Atomic):
    def __init__(self, name=None, parent=None, priority=None, rules=()):
        Atomic.__init__(self, name, parent, priority, config=rules)
        self.resources = Deque()
        self.store = Store()
        self.rules = None
        
    def status(self):
        return "%s: idle=%s\n%s" % (Atomic.status(self), 
                                    [req.origin.name for req in self.resources], 
                                    indent(self.store))
        
    def initializer(self):
        self.resources.clear()
        self.store.clear()
        self.rules = list(self.get_config())
        
    def receive(self, lot, _):
        self.store.add(lot)
        _, request, slots = self.check_resources()
        if slots:
            request.satisfy(self.get_rng().choice(slots))
            
    @chain
    def get_input(self, machine):
        request = Request(self.resources, machine)
        _, slots = self.check_resource(machine)
        if slots:
            request.satisfy(self.get_rng().choice(slots))
        yield request
        setup, proc, products = request.result
        yield Chain_result([setup, proc, self.store.get(products, verify=False)])
        
    def check_resources(self):
        choice = None
        slots = None
        max_score = None
        for request in self.resources:
            score, rsc_slots = self.check_resource(request.origin)
            if len(rsc_slots) > 0 and score > max_score:
                choice = request
                slots = rsc_slots
                max_score = score
        return max_score, choice, slots
        
    def check_resource(self, machine):
        slots = []
        max_score = [0.0] * len(self.rules)
        for setup, proc in machine.setup.slots():
            products = proc.check_availability(self.store.products())
            if products:
                score = [rule(machine, setup, proc, products) for rule in self.rules]
                if score > max_score:
                    max_score = score
                    slots = [(setup, proc, products)]
                elif score == max_score:
                    slots.append((setup, proc, products))
        return max_score, slots
        
class Line_controller(Atomic):
    def __init__(self, ID=None, parent=None, priority=None, rules=()):
        Atomic.__init__(self, ID, parent, priority, config=rules)
        self.pending_tasks = Request_list()
        self.idle_personnel = deque()
        self.rules = None
        
    def status(self):
        return "%s: idle=%s pending=%s" % (Atomic.status(self), 
                                           [person.get_name() for person in self.idle_personnel], 
                                           [req.origin.get_name() for req in self.pending_tasks])
                                           
    def initializer(self):
        self.pending_tasks.clear()
        self.idle_personnel.clear()
        self.rules = list(self.get_config())
        
    @chain
    def acquire_operator(self, machine):
        req = Request(self.pending_tasks, machine)
        operator = self.rsc4task(machine)
        if operator is not None:
            self.idle_personnel.remove(operator)
            req.satisfy(operator)
        yield req
        yield Chain_result(req.result)
        
    def idle(self, operator):
        req = self.task4rsc(operator)
        if req is not None:
            req.satisfy(operator)
        else:
            self.idle_personnel.append(operator)
            
    def rsc4task(self, task):
        if self.idle_personnel:
            n = len(self.idle_personnel)
            scores = [(self.score(rsc, task), n-i, rsc) \
                      for i, rsc in enumerate(self.idle_personnel)]
            return max(scores)[2]
        else:
            return None
            
    def task4rsc(self, rsc):
        if self.pending_tasks:
            n = len(self.pending_tasks)
            scores = [(self.score(rsc, req.data), n-i, req) \
                      for i, req in enumerate(self.pending_tasks)]
            return max(scores)[2]
        else:
            return None
            
    def score(self, resource, task):
        return [rule(resource, task) for rule in self.rules]

class Machine(Atomic):
    # state = Time_series_attr("state", timescale=HOUR, numeric=False, store=True)
    
    def __init__(self, ID=None, parent=None, priority=None, type=None, setup=None, 
                 dispatcher=None, controller=None, operation=None):
        Atomic.__init__(self, ID, parent, priority, config=setup)
        self.state = None # Machine.state.reset(self)
        self.dispatcher = dispatcher
        self.controller = controller
        self.current = None
        self.type = type
        self.setup = None
        self.operation = operation
        
    def status(self):
        return "%s: [%s-%s, %s] %s" % (Atomic.status(self), self.type, self.setup, 
                                       self.state, self.current)
        
    def initializer(self):
        self.state = "SB" # Machine.state.reset(self, "SB")
        self.setup = self.get_config()
        self.current = None
        return self.mainloop()
        
    @chain
    def mainloop(self):
        while True:
            setup, proc, input = (yield self.dispatcher.get_input(self)).result
            if setup is not self.setup:
                yield self.convert_to(setup)
            yield self.load(input)
            yield self.process(proc)
            yield self.unload()
            
    @chain
    def convert_to(self, setup):
        operator = (yield self.controller.acquire_operator(self)).result
        self.state = "SDS"
        yield operator.do("convert", self, setup)
        self.state = "SB"
        
    @chain
    def load(self, input):
        operator = (yield self.controller.acquire_operator(self)).result
        self.state = "PR"
        yield operator.do("load", self, input)
        
    @chain
    def process(self, proc):
        proc = proc.apply(self.current)
        tt = proc.tool_time()
        tbf, tol = proc.generate_fail()
        
        while tt > tbf:
            yield tbf
            yield self.recover(tol)
            tt -= tbf
            tbf, tol = proc.generate_fail()
        yield tt
        
        input, output = proc(self.get_root().product_table)
        assert input == []
        self.current = output
        
    @chain
    def unload(self):
        operator = (yield self.controller.acquire_operator(self)).result
        yield operator.do("unload", self)
        self.state = "SB"
        
    @chain
    def recover(self, tol):
        self.state = "UDW"
        operator = (yield self.controller.acquire_operator(self)).result
        self.state = "UDR"
        yield operator.do("recover", self, tol)
        self.state = "PR"
        
class Worker(Atomic):
    def __init__(self, ID=None, parent=None, priority=None, skillset=None, controller=None):
        Atomic.__init__(self, ID, parent, priority)
        self.skillset = skillset
        self.controller = controller
        self.current = None
        self.state = None
        
    def status(self):
        return "%s: [%s, %s] %s" % (Atomic.status(self), self.skillset, self.state, self.current)
        
    def initializer(self):
        self.current = None
        self.state = "idle"
        self.controller.idle(self)
        
    @contextmanager
    def __doing(self, s):
        self.state = s
        yield
        self.state = "idle"
        self.controller.idle(self)
        
    @bound_chain
    def do(self, skill, *args, **kwargs):
        with self.__doing(skill):
            yield self.skillset[skill](self, *args, **kwargs)
            
class Operation(Composite):
    def __init__(self, ID=None, parent=None, priority=None, 
                 dispatcher_rules=(), controller_rules=()):
        Composite.__init__(self, ID, parent, [], priority, None)
        self.equipment = Composite("equipment", self, priority=0.0)
        self.personnel = Composite("personnel", self, priority=1.0)
        self.controller = Line_controller("controller", self, priority=2.0, rules=controller_rules)
        self.dispatcher = Lot_dispatcher("dispatcher", self, priority=3.0, rules=dispatcher_rules)
        
        self.receive = self.dispatcher.receive
        self.personnel.add = self.addPersonnel
        self.equipment.add = self.addEquipment
        
    def add_personnel(self, *people):
        Composite.add(self.personnel, *people)
        for person in people:
            person.controller = self.controller
            
    def add_equipment(self, *machines):
        Composite.add(self.equipment, *machines)
        for machine in machines:
            machine.dispatcher = self.dispatcher
            machine.controller = self.controller
            machine.operation = self
            
class Job_shop(Composite):
    mandatory_config = ["products", # [Product]
                        "routes", # [Route]
                        "releases", # [Release]
                        #"WIP", # {opname: {prodname: [(size, no)]}}
                        "equipment_def", # [EqType]
                        "equipment_count", # {opname: {eqname: {setupname: count}}}
                        "personnel_def", # ?[Skillset]?
                        "personnel_count" # {opname: ?{skillsetname:? count}}
                        ]
    def __init__(self, ID=None, config=None):
        Composite.__init__(self, ID, config=config)
        self.linked = linked
        self.product_table = None
        
    def initializer(self):
        config = self.get_config()
        self.product_table = dict((p.name, p) for p in config.products)
        if not self.linked:
            self.link()
            self.linked = True
            
    def build(self):
        """Build the job shop."""
        for opname in self.operations:
            Operation(opname, self)
        
    def link_products(self):
        """Link products to respective routes."""
        route_index = dict((r.name, r) for r in self.get_config().routes)
        for product in config.products:
            product.route = route_index[product.route]
            
    def link_routes(self):
        """Link route sequences to respective operations (only possible after building)."""
        routes = set(p.route for p in self.get_config().products)
        for route in routes:
            linked_sequence = []
            for operation in route:
                linked_sequence.append(self[operation])
            route.sequence = linked_sequence
            
    def link_releases(self):
        """Link route sequences to respective operations (only possible after building)."""
        routes = set(p.route for p in self.get_config().products)
        for route in routes:
            linked_sequence = []
            for operation in route:
                linked_sequence.append(self[operation])
            route.sequence = linked_sequence
            
    def build_wip(self):
        """Prepare the correct WIP after building the model."""
        
    def build_releases(self):
        """Build a list of Release objects and create a LotRelease."""
        
    def build_equipment(self):
        """Build equipment according to the specified distribution."""
        
    def build_personnel(self):
        """Build personnel according to the specified distribution."""
        
# ++++++++++++++++++++++++++++++++++++++++++++++++
# WORKER SKILLS
# Skills MUST be chains (bound or unbound) which take two mandatory arguments and any number of 
# additional arguments. The two mandatory arguments are the worker and the duration of the task.
@chain
def convert(worker, duration, machine, setup):
    worker.current = machine
    yield duration + machine.setup.conversion_time(setup)
    worker.current = None
    machine.setup = setup
    
@chain
def load(worker, duration, machine, input):
    worker.current = input
    yield duration
    worker.current = None
    machine.current = input
    
@chain
def unload(worker, duration, machine):
    worker.current = machine
    yield duration
    worker.current = None
    output = machine.current
    machine.current = None
    for lot in output:
        try:
            target = lot.product.route.after(machine.operation)
        except (IndexError, ValueError):
            target = worker.find(">exit")
        worker.send(lot, target)
        
@chain
def recover(worker, duration, machine, tol):
    worker.current = machine
    yield duration + tol
    worker.current = None
    
skill_funcs = dict((f.__name__, f) for f in (convert, load, unload, recover))